package com.raos.example.webapi;

import com.raos.example.memory.entity.MyChatMessage;
import com.raos.example.memory.repository.MyChatMessageRepository;
import com.raos.example.util.R;
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.data.message.SystemMessage;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.model.chat.StreamingChatLanguageModel;
import dev.langchain4j.model.chat.response.ChatResponse;
import dev.langchain4j.model.chat.response.StreamingChatResponseHandler;
import jakarta.annotation.Resource;
import jakarta.servlet.http.HttpServletResponse;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.*;
import reactor.core.publisher.Flux;

import java.util.ArrayList;
import java.util.List;
import java.util.Objects;

/**
 * 开放聊天API
 *
 * @author raos
 * @date 2025-07-13 13:46
 */
@CrossOrigin
@Slf4j
@RestController
@RequestMapping("/openapi/chat")
public class OpenApiChat {

    @Resource(name = "ollamaStreamingChatModel")
    private StreamingChatLanguageModel streamingChatLanguageModel;
    @Autowired
    private MyChatMessageRepository chatMessageRepository;

    @PostMapping("/chat")
    public Flux<String> chat(@RequestParam(value = "sessionId") String sessionId,
                             @RequestParam(value = "text") String text,
                             HttpServletResponse response) {

        log.info("sessionId=[{}], text=[{}]", sessionId, text);

        List<MyChatMessage> messageList = chatMessageRepository.findAllByChatMemoryId(sessionId);

        List<ChatMessage> chatMessages = new ArrayList<>();

        // 添加数据库中的问答内容（本质就是上下文）
        for (MyChatMessage message1 : messageList) {
            ChatMessage chatMessage = message1.toChatMessage();
            if (Objects.nonNull(chatMessage)) {
                chatMessages.add(chatMessage);
            }
        }

        MyChatMessage myChatMessage = MyChatMessage.create(sessionId, "USER", text);
        myChatMessage = chatMessageRepository.save(myChatMessage);

        chatMessages.add(UserMessage.userMessage(myChatMessage.getContent()));

        // 解决响应乱码
        response.setContentType(MediaType.TEXT_EVENT_STREAM_VALUE + "; charset=utf-8");
        return Flux.create(fluxSink -> streamingChatLanguageModel.chat(chatMessages, new StreamingChatResponseHandler() {
            @Override
            public void onPartialResponse(String partialResponse) {
                // 每次响应的数据
                fluxSink.next(partialResponse);
            }

            @Override
            public void onCompleteResponse(ChatResponse completeResponse) {
                // 响应结束的回调
                log.info("***AI对话 onComplete:\n {}", completeResponse.aiMessage().text());
                MyChatMessage chatMessage = MyChatMessage.create(sessionId, "AI", completeResponse.aiMessage().text());
                chatMessageRepository.save(chatMessage);
                fluxSink.complete();
            }

            @Override
            public void onError(Throwable error) {
                // 响应出错的回调
                log.error("***AI对话，发送异常", error);
                fluxSink.error(error);
            }
        }));
    }

    /**
     * 获取历史会话列表
     */
    @PostMapping("/history")
    public R<List<MyChatMessage>> history(@RequestParam(value = "sessionId") String sessionId) {
        log.info("sessionId=[{}]", sessionId);
        List<MyChatMessage> messageList = chatMessageRepository.findAllByChatMemoryId(sessionId);
        return R.ok(messageList);
    }

}
